﻿#pragma once

#include "OCMain/Log.h"
#include "OCMain/ZipUtil.h"
#include "OCMain/FileUtil.h"
#include "OCVolume/VolumeHeader.h"
#include "OCGIS/ShpUtil.h"
#include "OCGIS/GDALCPP.h"

#define ClipZ 2.0
#define FrontSpeedClipFactor 0.6
#define EndSpeedClipFactor 0.9

//#define FilterMin 20
//#define SimplifyNum 2
//#define MaxTakenNum 2
//#define ZScale 0.05

#define FilterMin 25
#define SimplifyNum 18
#define MaxTakenNum 2
#define ZScale 1.0


using namespace OCPP;
using namespace OCPP::Volume;

point_type create_point(Vector3 point) {
	OGRPoint* p = (OGRPoint*)OGR_G_CreateGeometry(wkbPoint);
	p->setX(point.x);
	p->setY(point.y);
	p->setZ(point.z);
	return point_type{ p };
}

linestring_type create_linestring(Vector3List points) {
	OGRLineString* geom = (OGRLineString*)OGR_G_CreateGeometry(wkbLineString);
	for (auto p : points)
	{
		geom->addPoint(p.x, p.y, p.z);
	}
	return linestring_type{ geom };
}

class GeojsonWriter
{
public:
	GeojsonWriter(String fileName)
	{
		String name, ext, path;
		StringUtil::splitFullFilename(fileName, name, ext, path);
		try {
			dataset = std::make_unique<Dataset>("GeoJSON", fileName);
			//layer = std::make_unique<Layer>(*dataset, name, wkbPoint);
			layer = std::make_unique<Layer>(*dataset, name, wkbLineString);
			layer->add_field("id", OFTInteger, 4);
			layer->add_field("angle", OFTInteger, 4);
		}
		catch (gdal_error e) {
			std::cout << String{ "creating feature in layer '" } + "' failed" << "\n";
		}
	}

	void add(Vector3 point, int id, int angle)
	{
		Feature feature{ *layer,create_point(point) };
		feature.set_field("id", id);
		feature.set_field("angle", angle);
		feature.add_to_layer();
	}

	void add(Vector3List line, int id, int angle)
	{
		Feature feature{ *layer,create_linestring(line) };
		feature.set_field("id", id);
		feature.set_field("angle", angle);
		feature.add_to_layer();
	}

	std::unique_ptr<Layer> layer;
	std::unique_ptr<Dataset> dataset;
};

typedef std::function<Vector3(int x, int y, int z)> GetWindUVFunc;

class StreamGrid
{
public:
	StreamGrid(int width1, int height1, int depth1)
		:width(width1)
		, height(height1)
	{
		gridData.resize(depth1);
		for (auto& g : gridData)
		{
			g.resize(width * height);
		}
	}
	bool isTaken(int x, int y, int z, bool ocupy)
	{
		auto& grid = gridData[z];
		int i = y * width + x;
		if (grid[i] > MaxTakenNum)
		{
			return true;
		}
		if (ocupy)
		{
			grid[i] += 1;
		}
		return false;
	}

	int width, height;
	std::vector<Uint8Vector> gridData;
};

class StreamlineTracer
{
public:
	StreamlineTracer(GetWindUVFunc _getWind,
		int w, int h, int d = 1)
		:width(w),
		height(h),
		depth(d),
		getWind(_getWind),
		grid(w, h, d)
	{
		boundScale = { (float)width,(float)height,(float)depth };
	}

	Vector3 mix(Vector3 a, Vector3 b, float ratio) {
		return {
			a.x * ratio + (1 - ratio) * b.x,
			a.y * ratio + (1 - ratio) * b.y,
			a.z * ratio + (1 - ratio) * b.z
		};
	}

	Vector3 normalizedVectorField(Vector3 pos)
	{
		int lx = Math::Floor(pos.x);
		int ly = Math::Floor(pos.y);
		int lz = Math::Floor(pos.z);
		int ux = Math::Ceil(pos.x);
		int uy = Math::Ceil(pos.y);
		int uz = Math::Ceil(pos.z);

		if (lx < 0) lx = ux;
		if (ux >= width) ux = lx;
		if (ly < 0) ly = uy;
		if (uy > height) uy = ly;
		if (lz < 0) lz = uz;
		if (uz >= depth) uz = lz;

		if (lx < 0 || lx >= width ||
			ly < 0 || ly >= height ||
			ly < 0 || lz >= depth ||
			uz < 0 || ux >= width ||
			uy < 0 || uy >= height ||
			uz < 0 || uz >= depth)
		{
			return Vector3::ZERO;
		}

		Vector3 tl = getWind(lx, ly, lz);
		Vector3 tr = getWind(lx + 1, ly, lz);
		Vector3 bl = getWind(lx, ly + 1, lz);
		Vector3 br = getWind(lx + 1, ly + 1, lz);

		if (tl == Vector3::ZERO ||
			tr == Vector3::ZERO ||
			bl == Vector3::ZERO ||
			br == Vector3::ZERO)
		{
			return Vector3::ZERO;
		}

		Vector3 res = mix(mix(tl, tr, pos.x - lx), mix(bl, br, pos.x - lx), 1 - pos.y + ly);

		return res.normalisedCopy();
	}

	Vector3 rk4(Vector3 point)
	{
		point *= boundScale;

		Vector3 tl = normalizedVectorField(point);

#if 1
		Vector3 tr = normalizedVectorField(point + Vector3(1, 0, 0));
		Vector3 bl = normalizedVectorField(point + Vector3(0, 1, 0));
		Vector3 br = normalizedVectorField(point + Vector3(1, 1, 0));

		Vector2 f;
		f.x = point.x - Math::Floor(point.x);
		f.y = point.y - Math::Floor(point.y);

		Vector3 res = mix(mix(tl, tr, f.x), mix(bl, br, f.x), f.y);
		tl = res;
#endif 
		tl.z *= ZScale;

		return tl / boundScale * timeStep;
	}

	bool isTaken(Vector3 candidate, bool ocupy = true)
	{
		if (outOfBound(candidate))
		{
			return true;
		}
		candidate *= boundScale;
		return grid.isTaken(candidate.x, candidate.y, candidate.z, ocupy);
	}

	bool outOfBound(Vector3 pos)
	{
		return pos.x < 0.0 ||
			pos.x >= 1.0 ||
			pos.y < 0.0 ||
			pos.y >= 1.0 ||
			pos.z < 0.0 ||
			pos.z >= 1.0;
	}

	Vector3 growByVelocity(Vector3 pos, Vector3 velocity) {

		Vector3 candidate = pos + velocity;
		if (candidate.z < 0.0)
		{
			candidate.z = 0.0;
		}

		if (outOfBound(candidate))
		{
			return Vector3::ZERO;
		}

		if (isTaken(candidate)) {
			return Vector3::ZERO;
		}

		return candidate;
	}

	Vector3 growBackward(Vector3 pos)
	{
		Vector3 velocity = rk4(pos);
		if (velocity == Vector3::ZERO)
		{
			return Vector3::ZERO;
		}
		return growByVelocity(pos, -velocity);
	}

	Vector3 growForward(Vector3 pos)
	{
		Vector3 velocity = rk4(pos);
		if (velocity == Vector3::ZERO)
		{
			return Vector3::ZERO;
		}
		return growByVelocity(pos, velocity);
	}

	Vector3 tracerLine(Vector3 start)
	{
		Vector3List line;
		Vector3 current = start;
		while (true)
		{
			if (outOfBound(current))
			{
				break;
			}

			Vector3 point = growForward(current);
			if (point != Vector3::ZERO)
			{
				current = point;
				line.push_back(point);
			}
			else
			{
				break;
			}
		}

		current = start;
		while (true)
		{
			if (outOfBound(current))
			{
				break;
			}

			Vector3 point = growBackward(current);
			if (point != Vector3::ZERO)
			{
				current = point;
				line.insert(line.begin(), point);
			}
			else
			{
				break;
			}
		}
		if (line.size() > 2)
		{
			lines.push_back(line);
			return line[0];
		}

		return Vector3::ZERO;
	}

	Vector3 getNextStart()
	{
		float nextTimeStep = MaxTakenNum;

		for (; lastLineIndex < lines.size(); lastLineIndex++, currentLineCheckedIndex = lines[lastLineIndex - 1].size() * 1.0)
		{
			auto& line = lines[lastLineIndex];

			for (auto i = currentLineCheckedIndex; i < line.size(); i++)
			{
				auto p = line[i];
				Vector3 velocity = normalizedVectorField(p * boundScale);
				if (velocity == Vector3::ZERO)
				{
					continue;
				}
				velocity /= boundScale;

				Vector3 nextStart;
				nextStart.x = p.x - velocity.x * nextTimeStep;
				nextStart.y = p.y + velocity.y * nextTimeStep;
				nextStart.z = p.z;


				if (!isTaken(nextStart, false) && nextStart != Vector3::ZERO
					&& !outOfBound(nextStart))
				{
					return nextStart;
				}

				nextStart.x = p.x + velocity.x * nextTimeStep;
				nextStart.y = p.y - velocity.y * nextTimeStep;
				nextStart.z = p.z;


				if (!isTaken(nextStart, false) && nextStart != Vector3::ZERO
					&& !outOfBound(nextStart))
				{
					return nextStart;
				}

				nextStart.x = p.x + velocity.x * nextTimeStep;
				nextStart.y = p.y + velocity.y * nextTimeStep;
				nextStart.z = p.z;

				if (!isTaken(nextStart, false) && nextStart != Vector3::ZERO
					&& !outOfBound(nextStart))
				{
					return nextStart;
				}
			}
		}

		return Vector3::ZERO;
	}

	void run()
	{
		float size = 4.0;
		for (float r = 1.0; r < size; r++)
		{
			for (float c = 1.0; c < size; c++)
			{
				Vector3 start = { c / size,  r / size, 0.0 };

				while (true)
				{
					Vector3 validPoint = tracerLine(start);
					if (validPoint == Vector3::ZERO)
					{
						currentLineCheckedIndex++;
					}
					start = getNextStart();
					if (start == Vector3::ZERO)
					{
						break;
					}
				}
			}
		}

	}

	Vector3 boundScale;
	size_t lastLineIndex = 0;
	size_t currentLineCheckedIndex = 0;
	StreamGrid grid;
	std::vector<Vector3List> lines;
	float timeStep = 0.5;
	GetWindUVFunc getWind;
	int width, height, depth = 1;
};

void calcStreamline(String fileName);

void testStreamline()
{
	//String fileName = "E:/test/u_v-component_of_wind_isobaric_20230517030000.bin.gz";
	//calcStreamline(fileName);

	String path1 = "E:/test/wind3d/20220914";
	String path = "E:/test/2024/Vertical";
	StringVector files = FileUtil::findFiles(path);
	for (const auto& fileName : files)
	{
		if (FileUtil::getFileExtension(fileName) == "gz")
		{
			calcStreamline(joinPath(path, fileName));
		}
	}
};

void calcStreamline(String fileName)
{
	OCPP_LOG_INFO << "testStreamline" << "\n";
	//String fileName = "E:/test/u-component_f_wind_iobaric_v-component_f_wind_iobaric_Vertical_velocity_pressure_isobaric_20220914000000.bin.gz";

	String content = FileUtil::getFileBinary(fileName);
	String raw = ZipUtil::gzipDecompress(content);
	w3dmheader header;
	const char* ptr = raw.data();
	memcpy(&header, ptr, sizeof(w3dmheader));
	ptr += sizeof(w3dmheader);

	String featureJson;
	featureJson.resize(header.featureTableJSONByteLength);
	memcpy((char*)featureJson.c_str(), ptr, featureJson.size());
	ptr += featureJson.size();

	Json::Value value;
	Json::Reader reader;
	reader.parse(featureJson, value);

	float west = value["rectangle"]["west"].asFloat();
	float south = value["rectangle"]["south"].asFloat();
	float east = value["rectangle"]["east"].asFloat();
	float north = value["rectangle"]["north"].asFloat();

	int width = value["source"]["width"].asInt();
	int height = value["source"]["height"].asInt();
	int depth = value["source"]["depth"].asInt();

	int pixelFormat = value["source"]["pixelFormat"].asInt();

	int component = pixelFormat == GL_RGB ? 3 : pixelFormat == GL_RG ? 2 : 1;

	int size = width * height * depth;
	FloatVector floatData;
	floatData.resize(size * component);

	uint32 byteLength = floatData.size() * sizeof(float);

	assert(byteLength == header.featureTableBinaryByteLength);

	memcpy((char*)&floatData[0], ptr, byteLength);

	float dx = (east - west) / (width - 1);
	float dy = (north - south) / (height - 1);

	auto getUV = [&width, &height, &depth, &component, &floatData](int x, int y, int z) {
		if (z >= depth * ClipZ || x >= width || y >= height) return Vector3::ZERO;
		int i = (z * width * height + y * width + x) * component;
		float u = floatData[i];
		float v = floatData[i + 1];
		float w = component > 2 ? floatData[i + 2] : std::max(std::abs(u), std::abs(v));
		return Vector3(u, v, w);
		};

	StreamlineTracer tracer(getUV, width, height, depth);

	int id = 0;

#if 0
	GeojsonWriter writer("E:/test/streamline.geojson");
	for (int r = 0; r < height; r++)
	{
		float latitude = south + r * dy;
		for (int c = 0; c < width; c++)
		{
			float longitude = west + c * dx;
			int i = (r * width + c) * component;
			float u = floatData[i];
			float v = floatData[i + 1];
			Vector2 dir(u, v);
			auto angle = dir.angleBetween(Vector2(0, 1)).valueDegrees();
			if (u < 0.0) angle = 360.0 - angle;
			writer.add({ longitude,latitude,0.0 }, id++, angle);
		}
}
#else
	auto trimBySpeed = [&getUV, &width, &height, &depth](Vector3List& line, float speed)
		{
			int32 beg = 0;
			int32 end = 0;

			for (const auto& point : line)
			{
				auto uvw = getUV(point.x * width, point.y * height, point.z * depth);
				Vector2 uv(uvw.x, uvw.y);
				if (uv.length() > speed * FrontSpeedClipFactor)
				{
					break;
				}
				beg++;
			}

			for (auto it = line.rbegin(); it != line.rend(); ++it) {
				auto point = *it;
				auto uvw = getUV(point.x * width, point.y * height, point.z * depth);
				Vector2 uv(uvw.x, uvw.y);
				if (uv.length() > speed * EndSpeedClipFactor)
				{
					break;
				}
				end++;
			}

			end = (int32)line.size() - end;

			return std::make_tuple(beg, end);
		};

	String name, ext, path;
	StringUtil::splitFullFilename(fileName, name, ext, path);

	String outFileName = "\\\\10.104.207.155\\Cesium\\Assets\\2023\\Sample\\Wind3D\\20220914_850_300\\" + name + ".geojson";
	//String outFileName = "E:/test/streamline3.geojson";
	GeojsonWriter writer(outFileName);
	tracer.run();
	size_t writeCount = 0;
	for (auto& line : tracer.lines)
	{
		Vector3List linePoints;
		size_t count = 0;
		size_t simplify = SimplifyNum;
		float minAngle = 360.0;
		float maxAngle = 0.0;
		auto [beg, end] = trimBySpeed(line, FilterMin);
		if (end - beg > SimplifyNum)
		{
			for (; beg != end; beg++)
			{
				if (simplify > 0)
				{
					if (count % simplify != simplify - 1)
					{
						count++;
						continue;
					}
					count = 0;
				}

				const auto& point = line[beg];

				float longitude = west + point.x * (east - west);
				float latitude = south + point.y * (north - south);
				Vector3 p = { longitude ,latitude, point.z };
				linePoints.push_back(p);
			}
		}

		if (linePoints.size() > 1)
		{
			writer.add(linePoints, id++, 0.0);
			writeCount++;
		}
	}
#endif

	OCPP_LOG_INFO << featureJson << "\n";
	OCPP_LOG_INFO << writeCount << "\n";
}
